# -*- coding: utf-8 -*- 
"""
Project: AppleSnApi
Creator: SF
Create time: 2025-06-16 18:21
IDE: PyCharm
Introduction: worker 官网查询实现
"""
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
import random
from requests.exceptions import RequestException
import time
import json
import requests
from loguru import logger
import sys
import os
from datetime import datetime
from lxml import etree

current_path = os.path.abspath(os.path.dirname(__file__))
rootPath = os.path.split(current_path)[0]
sys.path.append(rootPath)

from settings import LOG_DIR, RANDOM_UA
from tools.helper import ddddOcr2captcha, get_system_info
from tools.DbConnect import redis_client, PROXY_MANAGER, sn_collection

# 配置日志
logger.add(
    os.path.join(LOG_DIR, 'spider.log'),
    rotation="500 MB",
    serialize=True,
    level='INFO',
    format="{time:YYYY-MM-DD HH:mm:ss} {level} From {module}.{function} : {message}",
    filter=lambda record: record["extra"].get("source") == "spider",
    encoding='utf-8'
)

spider_logger = logger.bind(source="spider")


# 格式化数据
def format_data(json_resp: str):
    """格式化官网返回数据"""
    tree = etree.HTML(json_resp)

    # 是否无效SN
    sorry_header = tree.xpath('//*[@class="sorry-header"]/span/text()')  # SN无效才有
    sorry_text = sorry_header[0].strip() if sorry_header else ''

    title_ele = tree.xpath('//*[@class="header-title"]/text()')  # 激活设备才有
    title = title_ele[0].strip() if title_ele else ''

    desc = tree.xpath('//*[@class="notification-description"]/text()')  # 未激活设备才有
    typeDesc = desc[0].strip() if desc else ''

    if sorry_text:
        if sorry_text == '无法验证序列号。请再试一次。':
            typeCode = '10003'
            typeMsg = '输入的序列号无效'
        else:
            typeCode = '10010'
            typeMsg = '未知错误'
        typeDesc = sorry_text
    else:
        if title and title == '有限保修':
            typeCode = '10001'
            typeMsg = '已激活'
            typeDesc = '设备已激活，查看获取支持'
        else:
            if '购买日期未验证' in json_resp or typeDesc == 'Apple 无法提供你的保障信息。请联系 Apple 支持以更新你的产品购买日期。':
                typeCode = '10002'
                typeMsg = '未激活'
            else:
                typeCode = '10010'
                typeMsg = '未知错误'

    validity_ele = tree.xpath('//*[@class="header-sub-title"]/*/text()')
    validity = validity_ele[0].strip() if validity_ele else ''

    product_ele = tree.xpath('//*[@class="content-area"]/div')
    if product_ele:
        img_ele = product_ele[0].xpath('.//*[@class="device-header-image"]/@src')  # 图片
        title_ele = product_ele[0].xpath('.//*[@class="device-header-title"]/text()')  # 机型
        sn_ele = product_ele[0].xpath('.//*[@class="device-header-serial-number"]/text()')  # 序列号
        purchase_ele = product_ele[0].xpath('.//*[@class="device-header-purchase"]/text()')  # 购买日期

        productInfo = {
            'imageUrl': img_ele[0].strip() if img_ele else '',
            'productConfig': '',
            'productName': '',
            'productType': title_ele[0].strip() if title_ele else '',
            'serialNumber': sn_ele[0].strip() if sn_ele else '',
            'dop': validity,
            'purchaseDate': purchase_ele[0].strip() if purchase_ele else '',
            'countryOfPurchase': '',
            'removeAirpodsLink': '',
            'productMeta': {
                'groupFamily': 'pgf31015',
                'superGroup': 'sg003',
                'prodFamilyClass': 'pfc3001',
                'name': title_ele[0].strip() if title_ele else ''
            }
        }
    else:
        productInfo = None

    return {
        'status': 'failed' if typeCode == '10010' else 'complete',
        'typeCode': typeCode,
        'typeMsg': typeMsg,
        'typeDesc': typeDesc,
        'productInfo': productInfo,
        'title': title if title else '',
        'validity': validity,
        'agreementNumber': '',
        'updated_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        'ss': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    }


# Apple官网查询
class AppleSNChecker:

    def __init__(self):
        self.manual_cookies = {'CCW_AFFINITY': 'awsapse1-2'}  # 手动传入
        self.session = requests.Session()
        self.session.headers.update({
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
            'Accept-Encoding': 'gzip, deflate, br',
            'Accept-Language': 'zh-CN,zh;q=0.9',
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36'
        })
        self.csrf_token = None
        self.cookies = {}  # 存储手动提取的cookie
        self.proxies = PROXY_MANAGER.get_proxy()  # 使用相同的代理
        self.timeout = 10
        # 初始化并执行前4步
        self._initialize(self.manual_cookies)

    def _initialize(self, manual_cookies=None):
        """执行初始化步骤(1-4)并设置必要的认证信息"""
        # 步骤1
        if not self.step_1_get_redirects_location():
            raise RuntimeError("步骤1失败: 无法重定向location")

        # 步骤2
        redirect_url = self.step_2_get_csrf_token()
        if not redirect_url:
            raise RuntimeError("步骤2失败: 无法获取csrf_token")

        # 步骤3
        if not self.step_3_page_init():
            raise RuntimeError("步骤3失败: 无法获取隐私授权同意")

        # 手动添加cookie（如果提供）
        if manual_cookies:
            for name, value in manual_cookies.items():
                self.session.cookies.set(name, value)
                self.cookies[name] = value

        # 打印最终的认证信息
        spider_logger.info(f"init cookie & csrf_token: {self.csrf_token}")

    # 提取cookie
    def extract_cookies(self, response):
        """从响应中提取所有cookie并存储"""
        for cookie in response.cookies:
            self.cookies[cookie.name] = cookie.value
        return self.cookies

    # 组合cookie为字符串
    def get_cookie_header(self):
        """生成Cookie请求头"""
        return "; ".join([f"{k}={v}" for k, v in self.cookies.items()])

    # 获取初始cookie
    def step_1_get_redirects_location(self):
        """步骤1: 获取重定向地址"""
        try:
            url = "https://checkcoverage.apple.com/?locale=zh_CN"
            response = self.session.get(url, proxies=self.proxies, timeout=self.timeout, verify=False,
                                        allow_redirects=False)
            response.raise_for_status()
            # 提取并存储cookie
            self.extract_cookies(response)

            # 从响应头中获取Location(重定向后的地址)
            location_url = response.headers.get('Location')
            if not location_url:
                raise ValueError("步骤1: 获取重定向地址失败")
            return True
        except RequestException:
            return False
        except ValueError:
            return False
        except Exception:
            return False

    # 重定向后获取csrf_token
    def step_2_get_csrf_token(self):
        """步骤2: 获取csrf_token"""
        try:
            url = "https://checkcoverage.apple.com/user-consent?locale=zh_CN"

            # 设置手动管理的cookie
            headers = {
                'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
                'Accept-Encoding': 'gzip, deflate, br',
                'Accept-Language': 'zh-CN,zh;q=0.9',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36',
                'Cookie': self.get_cookie_header()
            }

            response = self.session.get(url, headers=headers, proxies=self.proxies, timeout=self.timeout, verify=False)
            response.raise_for_status()

            # 更新cookie
            self.extract_cookies(response)

            data = response.text
            tree = etree.HTML(data)  # xpath 提取
            csrf_token = tree.xpath("//*[@name='csrf-token']/@content")
            if not csrf_token:
                raise ValueError("步骤2: 未找到csrf_token")
            self.csrf_token = csrf_token[0].strip()
            return True
        except RequestException:
            return False
        except (json.JSONDecodeError, ValueError):
            return False
        except Exception:
            return False

    # 对应页面获取隐私授权
    def step_3_page_init(self):
        """步骤3: 验证csrf_token 是否可用"""
        if not self.csrf_token:
            return False

        try:
            # 设置手动管理的cookie
            headers = {
                'Accept-Encoding': 'gzip, deflate, br',
                'Accept-Language': 'zh-CN,zh;q=0.9',
                'Origin': 'https://checkcoverage.apple.com',
                'Referer': 'https://checkcoverage.apple.com/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36',
                'accept': 'application/json',
                'content-type': 'application/json',
                'X-Apple-Csrf-Token': self.csrf_token,
                'X-Apple-SSE-ONLINE-CLIENT-APP-NAME': 'CoverageCentralWeb',
                'Cookie': self.get_cookie_header()
            }

            url = "https://checkcoverage.apple.com/api/v1/consent?locale=zh_CN"
            response = self.session.post(url, headers=headers, proxies=self.proxies, timeout=self.timeout, verify=False)
            response.raise_for_status()
            # 更新cookie
            self.extract_cookies(response)

            data = response.json()
            userConsented = data.get('meta').get("status")
            if userConsented != 'SUCCESS':
                raise ValueError("步骤3: 未成功进行隐私授权同意")
            return True
        except RequestException:
            return False
        except (json.JSONDecodeError, ValueError):
            return False
        except Exception:
            return False

    # 获取验证码
    def get_captcha(self):
        """步骤4: 获取验证码（独立方法）"""
        if not self.csrf_token:
            return False

        try:
            # 设置手动管理的cookie
            headers = {
                'Accept-Encoding': 'gzip, deflate, br',
                'Accept-Language': 'zh-CN,zh;q=0.9',
                'Origin': 'https://checkcoverage.apple.com',
                'Referer': 'https://checkcoverage.apple.com/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36',
                'accept': 'application/json',
                'content-type': 'application/json',
                'X-Apple-Csrf-Token': self.csrf_token,
                'X-Apple-SSE-ONLINE-CLIENT-APP-NAME': 'CoverageCentralWeb',
                'Cookie': self.get_cookie_header()
            }

            url = "https://checkcoverage.apple.com/api/v1/captcha?locale=zh_CN"
            response = self.session.get(url, headers=headers, proxies=self.proxies, timeout=self.timeout, verify=False)
            response.raise_for_status()

            # 更新cookie
            self.extract_cookies(response)

            data = response.json()
            captcha_image = data.get('data').get("binary")
            if not captcha_image:
                raise ValueError("步骤4: 未获取到验证码图片")
            return {'image': captcha_image}
        except RequestException:
            return False
        except (json.JSONDecodeError, ValueError):
            return False
        except Exception:
            return False

    # 校验验证码
    def check_captcha(self, sn, captcha_text):
        """步骤5: 提交验证码和SN（独立方法）"""
        if not self.csrf_token:
            return False

        try:
            headers = {
                'Accept-Encoding': 'gzip, deflate, br',
                'Accept-Language': 'zh-CN,zh;q=0.9',
                'Origin': 'https://checkcoverage.apple.com',
                'Referer': 'https://checkcoverage.apple.com/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36',
                'accept': 'application/json',
                'content-type': 'application/json',
                'X-Apple-Csrf-Token': self.csrf_token,
                'X-Apple-SSE-ONLINE-CLIENT-APP-NAME': 'CoverageCentralWeb',
                'Cookie': self.get_cookie_header()
            }

            url = "https://checkcoverage.apple.com/api/v1/captchaValidate?locale=zh_CN"
            payload = {
                "serialInput": sn,
                "answer": captcha_text,
                "captchaType": "image"
            }

            response = self.session.post(url, json=payload, headers=headers, proxies=self.proxies, timeout=self.timeout,
                                         verify=False)
            response.raise_for_status()

            # 更新cookie
            self.extract_cookies(response)

            if response.status_code == 200:  # 验证码通过
                return True
            elif response.status_code == 400:
                return False
            else:
                return False
        except RequestException:
            return False
        except json.JSONDecodeError:
            return False
        except Exception:
            return False

    # 查询激活结果
    def check_coverage(self):
        """步骤6: 获取最终查询结果"""
        if not self.csrf_token:
            return False

        try:
            headers = {
                'Accept-Encoding': '*/*',
                'Accept-Language': 'zh-CN,zh;q=0.9',
                'Origin': 'https://checkcoverage.apple.com',
                'Referer': 'https://checkcoverage.apple.com/',
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36',
                'accept': 'application/json',
                'content-type': 'application/json',
                'Next-Router-State-Tree': '%5B%22%22%2C%7B%22children%22%3A%5B%22__PAGE__%3F%7B%5C%22locale%5C%22%3A%5C%22zh_CN%5C%22%7D%22%2C%7B%7D%2C%22%2F%3Flocale%3Dzh_CN%22%2C%22refresh%22%5D%7D%2Cnull%2Cnull%2Ctrue%5D',
                'Next-Url': '/',
                'Cookie': self.get_cookie_header()
            }

            _rsc = ['x5c21', '1n8m2']
            url = f"https://checkcoverage.apple.com/coverage?locale=zh_CN&_rsc={random.choice(_rsc)}"

            response = self.session.get(url, headers=headers, proxies=self.proxies, timeout=self.timeout, verify=False)
            response.raise_for_status()

            data = response.content.decode('utf-8')
            if "NEXT_REDIRECT;replace;/user-consent?locale=zh_CN;307;" in data:
                return False
            else:
                return data
        except RequestException:
            return False
        except json.JSONDecodeError:
            return False
        except Exception:
            return False


# 查询函数
def verify_sn(params: dict):
    try:
        checker = AppleSNChecker()
        captcha = checker.get_captcha()
        if not captcha:
            return {'params': params, 'json_resp': 'get captcha error'}
        captcha_content = captcha.get("image")
        ddddocr_captcha_text = ddddOcr2captcha(captcha_content)
        if not ddddocr_captcha_text or len(ddddocr_captcha_text) <= 3:  # 识别失败或不足4位
            return {'params': params, 'json_resp': 'ocr error'}

        # 校验captcha
        sn = params.get("sn")
        captchaValidate = checker.check_captcha(sn.upper(), ddddocr_captcha_text)
        if captchaValidate:
            # 获取查询结果
            result = checker.check_coverage()
            if result:
                return {'params': params, 'json_resp': result}
    except RuntimeError as e:  # 各步骤出错
        spider_logger.error(f'抓取中有步骤出错: {str(e)}')
        return {'params': params, 'json_resp': None}
    except Exception as e:
        spider_logger.error(f"抓取中发生未知错误: {str(e)}")
        return {'params': params, 'json_resp': None}
    return {'params': params, 'json_resp': None}


# 爬虫主函数
def SpiderMain(worker=10):
    while redis_client.exists(f'sn_query_tasks'):
        with ThreadPoolExecutor(max_workers=worker) as executor:
            obj_list = []
            for i in range(int(worker)):
                item = redis_client.lpop(f'sn_query_tasks')
                if item:
                    obj = executor.submit(verify_sn, json.loads(item))
                    obj_list.append(obj)

            for future in as_completed(obj_list):
                try:
                    rst = future.result()
                    params, json_resp = rst.get('params'), rst.get('json_resp')

                    if json_resp is None:
                        redis_client.rpush(f'sn_query_tasks', json.dumps(params, ensure_ascii=False))
                        redis_client.hincrby(f'apple_sn_counter', 'failed_count')
                        continue

                    if isinstance(json_resp, str):
                        if json_resp == 'get captcha error':  # 获取验证码出错
                            redis_client.hincrby(f'apple_sn_counter', 'get_captcha_error')
                            redis_client.rpush(f'sn_query_tasks', json.dumps(params, ensure_ascii=False))

                        elif json_resp == 'ocr error':  # OCR 识别出错
                            redis_client.hincrby(f'apple_sn_counter', 'ocr_error')
                            redis_client.rpush(f'sn_query_tasks', json.dumps(params, ensure_ascii=False))

                        else:
                            query_data = format_data(json_resp)
                            status = query_data.get("typeCode")
                            if status == '10003':
                                redis_client.hincrby(f'apple_sn_counter', 'valid_sn')
                            elif status == '10010':
                                redis_client.hincrby(f'apple_sn_counter', 'unknown_error')
                            else:
                                redis_client.hincrby(f'apple_sn_counter', 'success')

                            # 更新数据库
                            params.update(query_data)
                            sn_collection.update_one({"sn": params.get("sn")}, {"$set": params})
                    else:
                        redis_client.rpush(f'sn_query_tasks', json.dumps(params, ensure_ascii=False))
                        redis_client.hincrby(f'apple_sn_counter', 'unknown_resp')
                        spider_logger.info(f'{params} >> got unknown resp:{json_resp}')
                except Exception:
                    redis_client.rpush(f'sn_query_tasks', json.dumps(params, ensure_ascii=False))

        time.sleep(random.uniform(0.2, 1.2))


if __name__ == '__main__':
    while True:
        try:
            queue_length = redis_client.llen("sn_query_tasks")
            poll_interval = random.uniform(10, 30)
            if queue_length > 0:
                spider_logger.info(f"启动SN查询Worker,当前节点检测到待查询任务:{queue_length} 个")
                SpiderMain(5)
            else:
                spider_logger.debug(f"当前任务队列为空，主动休眠{poll_interval} 秒")
                time.sleep(poll_interval)
        except Exception as e:
            spider_logger.error(f"worker 处理发生异常: {str(e)}")
            time.sleep(60)  # 发生异常后休眠较长时间
        else:
            # 检查主机状态并记录
            current_system_info = get_system_info(target_process='worker.py')
            current_host = current_system_info.get("hostname")
            redis_client.hset('apple_sn_counter', f'worker:{current_host}',
                              json.dumps(current_system_info, ensure_ascii=False))
